import csv
import os
import xml.etree.ElementTree as ET
import mujoco
import mujoco.viewer as viewer
import time
from threading import Thread
from scipy.spatial.transform import Rotation
import numpy as np
import cv2

def load_csv_to_dict_list(file_path):
    data = []
    with open(file_path, mode='r') as file:
        csv_reader = csv.DictReader(file)
        for row in csv_reader:
            data.append(row)
    return data


class RunMujoco:
    def __init__(
        self, 
        config,
        scene_dir,
        laser_path,
        resolution: float = 0.02,
        scan_range: list = [5, 5],
        output_dir: str = ""
    ):
        self.config = config
        self.laser_pos = [float(self.config["X"]), float(self.config["Y"]), 0.5]
        self.angle = float(self.config["RotationAngle"])
        self.image_name = self.config["Image"]
        
        self.scene_dir = scene_dir
        self.laser_path = laser_path
        
        self.model = mujoco.MjModel.from_xml_string(
            self.process_mjcf_str()
        )
        self.data = mujoco.MjData(self.model)
        
        self.done = False
        
        self.resolution = resolution
        self.scan_range = scan_range
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)
        
        self.laser_sensor_names = []
        for i in range(self.model.nsensor):
            sensor_name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_SENSOR, i)
            if "laser" in sensor_name:
                self.laser_sensor_names.append(sensor_name)
        self.laser_sensor_names = sorted(
            self.laser_sensor_names, 
            key=lambda x: int(x.replace("laser", ""))
        )
        self.launch()
        
    
    def process_mjcf_str(self):
        scene_name = os.path.basename(self.scene_dir)
        scene_path = os.path.join(self.scene_dir, f"{scene_name}.xml")
        with open(scene_path, 'r', encoding='utf-8') as file:
            xml_str = file.read()
        root = ET.fromstring(xml_str)
        include_tag = ET.Element("include", attrib={"file": self.laser_path})
        root.append(include_tag)
        modified_xml_str = ET.tostring(root, encoding='utf-8', method='xml').decode('utf-8')
        return modified_xml_str
        
    # 启动渲染
    def launch(self):
        self.run_viewer()
    
    # 启动GUI
    def run_viewer(self):
        self.set_laser_pos()
        mujoco.mj_step(self.model, self.data)
        self.get_laser_img()
        done = True
        
        
    # 设置激光雷达位置
    def set_laser_pos(self):
        self.data.qpos[: 3] = self.laser_pos
        # [x, y, z, w]
        quat_scipy = Rotation.from_euler(
            'xyz', 
            [0, 0, np.deg2rad(self.angle)]
        ).as_quat() 
        quat = [quat_scipy[3], quat_scipy[0], quat_scipy[1], quat_scipy[2]]
        self.data.qpos[3: 7] = quat
        
    def get_laser_img(self):
        laser_sensor_data = []
        for sensor in self.laser_sensor_names:
            laser_sensor_data.append(float(self.data.sensor(sensor).data[0]))
        cutoff = 10.0
        
        
        point_cloud = []
        for angle, data in enumerate(laser_sensor_data):
            if data > cutoff:
                continue
            x = float(data * np.cos(np.deg2rad(angle)))
            y = float(data * np.sin(np.deg2rad(angle)))
            if (
                x > 0 and
                x < self.scan_range[0] and 
                abs(y) < self.scan_range[1] / 2
            ):
                point_cloud.append([x, y])
        scan_img_size = (
            int(self.scan_range[0] / self.resolution),
            int(self.scan_range[1] / self.resolution)
        )
        gray_img = np.ones(scan_img_size, dtype=np.uint8) * 255
        for point in point_cloud:
            
            pixel_y = int((-point[1] + self.scan_range[1] / 2) / self.resolution)
            pixel_x = int(point[0] / self.resolution)
            if(
                pixel_y < 0 or pixel_y > scan_img_size[1] - 1 or
                pixel_x < 0 or pixel_x > scan_img_size[0] - 1
            ):
                continue
            center = (pixel_x, pixel_y)
            cv2.circle(gray_img, center, 2, 0, -1)
        result = cv2.rotate(gray_img, cv2.ROTATE_90_COUNTERCLOCKWISE)
        _, result = cv2.threshold(result, 127, 255, cv2.THRESH_BINARY)
        output_img = os.path.join(self.output_dir, self.image_name)
        cv2.imwrite(output_img, result)

# <include file="../../laser/laser.xml"/>
if __name__ == "__main__":
    
    image_data_csv = "dataset/image_data.csv"  # 替换为你的 CSV 文件路径
    image_data = load_csv_to_dict_list(image_data_csv)

    laser_path = "dataset_tools/sources/laser/laser.xml"
    scenes_dir = "dataset_tools/sources/mjcf_scenes"
        
    from concurrent.futures import ThreadPoolExecutor
    from tqdm import tqdm
    def process(data):
        scene_dir = os.path.join(scenes_dir, data["Scene"])
        runMujoco = RunMujoco(
            data,
            scene_dir,
            laser_path, 
            output_dir = "dataset/sample"
        )
        return runMujoco

    done_count = 0
    worker = 16
    with ThreadPoolExecutor(max_workers=worker) as executor:
        # results = executor.map(process, image_data)
        results = list(tqdm(executor.map(process, image_data), total=len(image_data), desc="Processing"))
        # done_count += 16
        # print(f"done: {done_count} / {len(image_data)}")